package org.movee.spring.data.hugegraph.schema;

import com.baidu.hugegraph.driver.HugeClient;
import com.baidu.hugegraph.driver.SchemaManager;
import com.baidu.hugegraph.structure.schema.EdgeLabel;
import com.baidu.hugegraph.structure.schema.IndexLabel;
import com.baidu.hugegraph.structure.schema.PropertyKey;
import com.baidu.hugegraph.structure.schema.VertexLabel;
import org.movee.spring.data.hugegraph.annotation.Edge;
import org.movee.spring.data.hugegraph.annotation.Index;
import org.movee.spring.data.hugegraph.annotation.NonNullable;
import org.movee.spring.data.hugegraph.annotation.Nullable;
import org.movee.spring.data.hugegraph.annotation.PrimaryKey;
import org.movee.spring.data.hugegraph.annotation.SortKey;
import org.movee.spring.data.hugegraph.annotation.Unique;
import org.movee.spring.data.hugegraph.annotation.Vertex;
import org.movee.spring.data.hugegraph.exception.HugeSchemaException;
import org.movee.spring.data.hugegraph.mapping.HugeEntityType;
import org.movee.spring.data.hugegraph.mapping.HugeMappingContext;
import org.movee.spring.data.hugegraph.mapping.HugePersistentEntity;
import org.movee.spring.data.hugegraph.mapping.HugePersistentProperty;
import com.google.common.base.Joiner;
import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.exception.ExceptionUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;

import java.util.ArrayList;
import java.util.Date;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.UUID;
import java.util.stream.Collectors;

@Slf4j
public class SchemaInitiator {

    private final Map<String, VertexLabel> vertexLabels = new HashMap<>();
    private final Map<String, EdgeLabel> edgeLabels = new HashMap<>();
    private Set<String> propertyKeys = new HashSet<>();
    private Set<String> indexLabels = new HashSet<>();

    private final List<HugePersistentEntity<?>> delayedEntities = new ArrayList<>();

    private final HugeClient hugeClient;
    private final SchemaManager schema;
    private final HugeMappingContext context;

    private volatile boolean initialized = false;


    public SchemaInitiator(final HugeClient hugeClient, final HugeMappingContext context) {
        this.hugeClient = hugeClient;
        this.context = context;
        this.schema = hugeClient.schema();
    }

    public void updateSchema(final HugePersistentEntity<?> entity) {

        if (hugeClient == null || entity == null) {
            log.error("to update schema, graph client and  persistent entity must not be null");
            return;
        }

        if (StringUtils.isEmpty(entity.getLabel())) {
            log.error("label value of persistent entity is null or empty, entity class: {}", entity.getName());
            return;
        }

        if (entity.getEntityType() != HugeEntityType.VERTEX
                && entity.getEntityType() != HugeEntityType.EDGE) {
            log.error("entity type not supported, entity class: {}, entity type: {}",
                    entity.getName(), entity.getEntityType());
            return;
        }

        try {
            // 1. 获取已创建的schema
            maybeGetStoredSchema();

            // 2. 创建edge需先创建相应的vertex，否则延时创建
            boolean delayed = maybeDelayEdgeEntity(entity);
            if (delayed) {
                return;
            }

            // 3. 创建schema
            maybeUpdateSchema(entity);

            // 4. edge的vertex已创建，创建对应的schema
            updateDelayedEntities(entity);
        } catch (Exception e) {
            log.error("update schema for {} failed, {}", entity.getName(), ExceptionUtils.getFullStackTrace(e));
        }

    }

    private boolean maybeDelayEdgeEntity(HugePersistentEntity<?> entity) {
        if (entity.getEntityType() == HugeEntityType.EDGE) {
            HugePersistentProperty headProperty = entity.getEdgeHeadProperty();
            HugePersistentProperty tailProperty = entity.getEdgeTailProperty();
            if (headProperty != null && tailProperty != null) {
                HugePersistentEntity<?> headEntity = context.getPersistentEntity(headProperty.getType());
                HugePersistentEntity<?> tailEntity = context.getPersistentEntity(tailProperty.getType());
                if (headEntity != null && tailEntity != null) {
                    if (!vertexLabels.containsKey(headEntity.getLabel())
                            || !vertexLabels.containsKey(tailEntity.getLabel())) {
                        // 还没有添加对应的vertex
                        log.info("delay edgeLabel: {}", entity.getLabel());
                        delayedEntities.add(entity);
                        return true;
                    }
                } else {
                    return true;
                }
            } else {
                return true;
            }
        }

        return false;
    }

    private void updateDelayedEntities(HugePersistentEntity<?> entity) {
        if (entity.getEntityType() == HugeEntityType.VERTEX) {

            Iterator<HugePersistentEntity<?>> itor = delayedEntities.iterator();
            while (itor.hasNext()) {
                HugePersistentEntity<?> delayed = itor.next();
                HugePersistentProperty headProperty = delayed.getEdgeHeadProperty();
                HugePersistentProperty tailProperty = delayed.getEdgeTailProperty();
                HugePersistentEntity<?> headEntity = context.getPersistentEntity(headProperty.getType());
                HugePersistentEntity<?> tailEntity = context.getPersistentEntity(tailProperty.getType());
                Assert.notNull(headEntity, "headEntity would not be null");
                Assert.notNull(tailEntity, "headEntity would not be null");
                if (vertexLabels.containsKey(headEntity.getLabel())
                        && vertexLabels.containsKey(tailEntity.getLabel())) {
                    log.info("update delayed edgeLabel: {}", delayed.getLabel());
                    maybeUpdateSchema(delayed);
                    itor.remove();
                }
            }

        }
    }


    private void maybeGetStoredSchema() {

        if (!initialized) {
            synchronized (this) {
                if (!initialized) {

                    log.info("do init, get schema from graph database");

                    List<PropertyKey> storedPropertyKeys = schema.getPropertyKeys();
                    propertyKeys = storedPropertyKeys.stream().map(PropertyKey::name).collect(Collectors.toSet());

                    List<VertexLabel> storedVertexLabels = schema.getVertexLabels();
                    storedVertexLabels.forEach(v -> vertexLabels.put(v.name(), v));

                    List<EdgeLabel> storedEdgeLabels = schema.getEdgeLabels();
                    storedEdgeLabels.forEach(e -> edgeLabels.put(e.name(), e));

                    List<IndexLabel> storedIndexLabels = schema.getIndexLabels();
                    indexLabels = storedIndexLabels.stream().map(IndexLabel::name).collect(Collectors.toSet());

                    log.info("stored PropertyKey: {}", propertyKeys);
                    log.info("sotred VertexLabel: {}", storedVertexLabels.stream()
                            .map(VertexLabel::name).collect(Collectors.toList()));
                    log.info("stored EdgeLabel: {}", storedEdgeLabels.stream()
                            .map(EdgeLabel::name).collect(Collectors.toList()));
                    log.info("stored IndexLabel: {}", indexLabels);

                    initialized = true;
                }
            }
        }
    }


    private void maybeUpdateSchema(final HugePersistentEntity<?> entity) {

        log.debug("to update schema for {}", entity.getName());

        maybeUpdatePropertyKeys(entity);

        if (entity.getEntityType() == HugeEntityType.VERTEX) {
            maybeUpdateVertexLabel(entity);
        } else {
            maybeUpdateEdgeLabel(entity);
        }

        maybeUpdateIndexLabels(entity);
        maybeUpdateUniqueIndexLabels(entity);
    }

    private void maybeUpdatePropertyKeys(final HugePersistentEntity<?> entity) {
        for (HugePersistentProperty property : entity.getPersistentProperties()) {
            String name = property.getName();
            if (!propertyKeys.contains(name)) {
                if (!property.isIdProperty()
                        && !property.isEdgeHeadProperty()
                        && !property.isEdgeTailProperty()) {

                    log.debug("create new PropertyKey: {}", name);
                    boolean success = createPropertyKey(property);
                    if (success) {
                        propertyKeys.add(name);
                    }
                }
            }
        }
    }

    private void maybeUpdateVertexLabel(final HugePersistentEntity<?> entity) {

        Set<String> properties = entity.getPersistentProperties().stream()
                .filter(p -> !p.isIdProperty())
                .map(HugePersistentProperty::getName)
                .collect(Collectors.toSet());

        Set<String> nullableProperties = entity.getPersistentProperties().stream()
                .filter(p -> !p.isIdProperty())
                .filter(p -> p.findAnnotation(Nullable.class) != null)
                .map(HugePersistentProperty::getName)
                .collect(Collectors.toSet());

        Set<String> nonNullableProperties = entity.getPersistentProperties().stream()
                .filter(p -> !p.isIdProperty())
                .filter(p -> p.findAnnotation(NonNullable.class) != null)
                .map(HugePersistentProperty::getName)
                .collect(Collectors.toSet());

        Set<String> nullableKeys = new HashSet<>(properties);
        if (!nullableProperties.isEmpty()) {
            nullableKeys = nullableProperties;
        } else if (!nonNullableProperties.isEmpty()) {
            nullableKeys.removeAll(nonNullableProperties);
        }

        VertexLabel storedVertexLabel = vertexLabels.get(entity.getLabel());
        if (storedVertexLabel == null) {
            VertexLabel.Builder builder = schema.vertexLabel(entity.getLabel());
            Vertex vertexAnno = entity.findAnnotation(Vertex.class);
            Assert.notNull(vertexAnno, "must not be null");
            switch (vertexAnno.idStrategy()) {
                case AUTOMATIC: {
                    builder.useAutomaticId();
                }
                break;
                case CUSTOMIZE_STRING: {
                    builder.useCustomizeStringId();
                }
                break;
                case CUSTOMIZE_NUMBER: {
                    builder.useCustomizeNumberId();
                }
                break;
                case PIRMARY_KEY: {
                    builder.usePrimaryKeyId();

                    List<String> primaryKeys = entity.getPersistentProperties().stream()
                            .filter(p -> p.findAnnotation(PrimaryKey.class) != null)
                            .map(HugePersistentProperty::getName)
                            .collect(Collectors.toList());

                    if (!CollectionUtils.isEmpty(primaryKeys)) {
                        builder.primaryKeys(primaryKeys.toArray(new String[0]));
                    } else {
                        throw new HugeSchemaException("PIRMARY_KEY idStrategy must have primary keys");
                    }
                }
                break;
                default: {
                    // no
                }
            }

            builder.enableLabelIndex(vertexAnno.enableLabelIndex());

            builder.properties(properties.toArray(new String[0])).nullableKeys(nullableKeys.toArray(new String[0]));

            VertexLabel label = builder.ifNotExist().create();
            log.info("create new VertexLabel: {}", label.name());
            vertexLabels.put(label.name(), label);
        } else {
            Set<String> storedProperties = storedVertexLabel.properties();
            Set<String> storedNullableKeys = storedVertexLabel.nullableKeys();

            properties.removeAll(storedProperties);
            nullableKeys.removeAll(storedNullableKeys);
            if (!properties.isEmpty() || !nullableKeys.isEmpty()) {
                VertexLabel.Builder builder = schema.vertexLabel(entity.getLabel());
                if (!properties.isEmpty()) {
                    builder.properties(properties.toArray(new String[0]));
                }
                if (!nullableKeys.isEmpty()) {
                    builder.nullableKeys(nullableKeys.toArray(new String[0]));
                }

                VertexLabel label = builder.append();
                log.info("update new VertexLabel: {}", label.name());
                vertexLabels.put(label.name(), label);
            }
        }


    }

    private void maybeUpdateEdgeLabel(final HugePersistentEntity<?> entity) {

        Set<String> properties = entity.getPersistentProperties().stream()
                .filter(p -> !p.isIdProperty() && !p.isEdgeHeadProperty() && !p.isEdgeTailProperty())
                .map(HugePersistentProperty::getName)
                .collect(Collectors.toSet());

        Set<String> nullableProperties = entity.getPersistentProperties().stream()
                .filter(p -> !p.isIdProperty() && !p.isEdgeHeadProperty() && !p.isEdgeTailProperty())
                .filter(p -> p.findAnnotation(Nullable.class) != null)
                .map(HugePersistentProperty::getName)
                .collect(Collectors.toSet());

        Set<String> nonNullableProperties = entity.getPersistentProperties().stream()
                .filter(p -> !p.isIdProperty() && !p.isEdgeHeadProperty() && !p.isEdgeTailProperty())
                .filter(p -> p.findAnnotation(NonNullable.class) != null)
                .map(HugePersistentProperty::getName)
                .collect(Collectors.toSet());

        List<String> sortKeys = entity.getPersistentProperties().stream()
                .filter(p -> p.findAnnotation(SortKey.class) != null)
                .map(HugePersistentProperty::getName)
                .collect(Collectors.toList());

        Set<String> nullableKeys = new HashSet<>(properties);
        if (!nullableProperties.isEmpty()) {
            nullableKeys = nullableProperties;
        } else if (!nonNullableProperties.isEmpty()) {
            nullableKeys.removeAll(nonNullableProperties);
        }
        // sort key不能是nullable
        for (String sortKey : sortKeys) {
            nullableKeys.remove(sortKey);
        }

        EdgeLabel storedEdgeLabel = edgeLabels.get(entity.getLabel());
        if (storedEdgeLabel == null) {
            HugePersistentProperty headProperty = entity.getEdgeHeadProperty();
            HugePersistentProperty tailProperty = entity.getEdgeTailProperty();
            if (headProperty != null && tailProperty != null) {
                HugePersistentEntity<?> headEntity = context.getPersistentEntity(headProperty.getType());
                HugePersistentEntity<?> tailEntity = context.getPersistentEntity(tailProperty.getType());
                if (headEntity != null && tailEntity != null) {

                    EdgeLabel.Builder builder = schema.edgeLabel(entity.getLabel());
                    Edge edgeAnno = entity.findAnnotation(Edge.class);
                    Assert.notNull(edgeAnno, "must not be null");
                    if (edgeAnno.multiplicity() == Multiplicity.MULTI) {
                        builder.multiTimes();
                    } else {
                        builder.singleTime();
                    }

                    builder.link(headEntity.getLabel(), tailEntity.getLabel());

                    builder.properties(properties.toArray(new String[0]))
                            .sortKeys(sortKeys.toArray(new String[0]))
                            .nullableKeys(nullableKeys.toArray(new String[0]));
                    EdgeLabel label = builder.ifNotExist().create();
                    log.info("create new EdgeLabel: {}", label.name());
                    edgeLabels.put(label.name(), label);
                }
            }

        } else {
            Set<String> storedProperties = storedEdgeLabel.properties();
            Set<String> storedNullableKeys = storedEdgeLabel.nullableKeys();

            properties.removeAll(storedProperties);
            nullableKeys.removeAll(storedNullableKeys);
            if (!properties.isEmpty() || !nullableKeys.isEmpty()) {
                EdgeLabel.Builder builder = schema.edgeLabel(entity.getLabel());
                if (!properties.isEmpty()) {
                    log.info("properties: {}", properties);
                    builder.properties(properties.toArray(new String[0]));
                }
                if (!nullableKeys.isEmpty()) {
                    log.info("nullableKeys: {}", nullableKeys);
                    builder.nullableKeys(nullableKeys.toArray(new String[0]));
                }

                EdgeLabel label = builder.append();
                log.info("update EdgeLabel: {}", label.name());
                edgeLabels.put(label.name(), label);
            }
        }
    }

    private void maybeUpdateIndexLabels(HugePersistentEntity<?> entity) {

        Index[] annoIndexes = null;
        if (entity.getEntityType() == HugeEntityType.VERTEX) {
            Vertex vertexAnno = entity.findAnnotation(Vertex.class);
            Assert.notNull(vertexAnno, "must not be null");
            annoIndexes = vertexAnno.indexes();
        } else {
            Edge edgeAnno = entity.findAnnotation(Edge.class);
            Assert.notNull(edgeAnno, "must not be null");
            annoIndexes = edgeAnno.indexes();
        }

        for (int i = 0; i < annoIndexes.length; i++) {
            createIndexLabel(annoIndexes[i], "", entity);
        }

        for (HugePersistentProperty property : entity.getPersistentProperties()) {

            Index indexAnno = property.findAnnotation(Index.class);
            if (indexAnno != null) {
                createIndexLabel(indexAnno, property.getName(), entity);
            }
        }
    }

    private void createIndexLabel(Index index, String propertyName, HugePersistentEntity<?> entity) {

        List<String> fields = Lists.newArrayList(index.fields());
        if (fields.isEmpty()) {
            if (!StringUtils.isEmpty(propertyName)) {
                fields.add(propertyName);
            } else {
                return;
            }
        }

        String indexName = index.name();
        if (StringUtils.isEmpty(indexName)) {
            String prefix = StringUtils.uncapitalize(entity.getType().getSimpleName());
            List<String> capFields = fields.stream().map(StringUtils::capitalize).collect(Collectors.toList());
            String postfix = Joiner.on("And").join(capFields);
            indexName = prefix + "By" + postfix;
        }

        if (!indexLabels.contains(indexName)) {
            IndexLabel.Builder builder = schema.indexLabel(indexName);
            if (entity.getEntityType() == HugeEntityType.VERTEX) {
                builder.onV(entity.getLabel());
            } else {
                builder.onE(entity.getLabel());
            }

            builder.by(fields.toArray(new String[0]));
            switch (index.type()) {
                case SECONDARY: {
                    builder.secondary();
                }
                break;
                case RANGE: {
                    builder.range();
                }
                break;
                case SEARCH: {
                    builder.search();
                }
                break;
                case SHARD: {
                    builder.shard();
                }
                break;
                default: {
                    // no
                }
            }

            IndexLabel label = builder.ifNotExist().create();
            log.info("create new IndexLabel: {}", label.name());
            indexLabels.add(label.name());
        }

    }

    private void maybeUpdateUniqueIndexLabels(HugePersistentEntity<?> entity) {

        Unique[] annoIndexes = null;
        if (entity.getEntityType() == HugeEntityType.VERTEX) {
            Vertex vertexAnno = entity.findAnnotation(Vertex.class);
            Assert.notNull(vertexAnno, "must not be null");
            annoIndexes = vertexAnno.uniqueConstraint();
        } else {
            Edge edgeAnno = entity.findAnnotation(Edge.class);
            Assert.notNull(edgeAnno, "must not be null");
            annoIndexes = edgeAnno.uniqueConstraint();
        }

        for (int i = 0; i < annoIndexes.length; i++) {
            createUniqueIndexLabel(annoIndexes[i], "", entity);
        }

        for (HugePersistentProperty property : entity.getPersistentProperties()) {

            Unique indexAnno = property.findAnnotation(Unique.class);
            if (indexAnno != null) {
                createUniqueIndexLabel(indexAnno, property.getName(), entity);
            }
        }
    }

    private boolean createPropertyKey(HugePersistentProperty property) {

        PropertyKey.Builder builder = schema.propertyKey(property.getName());

        // datatype
        Class<?> type = property.getActualType();
        if (type == String.class) {
            builder.asText();
        } else if (type == Integer.class) {
            builder.asInt();
        } else if (type == Date.class) {
            builder.asDate();
        } else if (type == UUID.class) {
            builder.asUUID();
        } else if (type == Boolean.class) {
            builder.asBoolean();
        } else if (type == Byte.class) {
            builder.asByte();
        } else if (type == Byte[].class) {
            builder.asBlob();
        } else if (type == Double.class) {
            builder.asDouble();
        } else if (type == Float.class) {
            builder.asFloat();
        } else if (type == Long.class) {
            builder.asLong();
        } else {
            log.warn("unsupported property key. name: {}, type: {}", property.getName(), type.getName());
            return false;
        }

        // cardinality
        Class rawType = property.getRawType();
        if (rawType == List.class) {
            builder.valueList();
        } else if (rawType == Set.class) {
            builder.valueSet();
        } else {
            builder.valueSingle();
        }

        builder.ifNotExist().create();
        return true;
    }


    private void createUniqueIndexLabel(Unique index, String propertyName, HugePersistentEntity<?> entity) {


        List<String> fields = Lists.newArrayList(index.fields());
        if (fields.isEmpty()) {
            if (!StringUtils.isEmpty(propertyName)) {
                fields.add(propertyName);
            } else {
                return;
            }
        }

        String indexName = index.name();
        if (StringUtils.isEmpty(indexName)) {
            String prefix = StringUtils.uncapitalize(entity.getType().getSimpleName());
            List<String> capFields = fields.stream().map(StringUtils::capitalize).collect(Collectors.toList());
            String postfix = Joiner.on("And").join(capFields);
            indexName = prefix + "UniqueBy" + postfix;
        }

        if (!indexLabels.contains(indexName)) {
            IndexLabel.Builder builder = schema.indexLabel(indexName);
            if (entity.getEntityType() == HugeEntityType.VERTEX) {
                builder.onV(entity.getLabel());
            } else {
                builder.onE(entity.getLabel());
            }

            builder.by(fields.toArray(new String[0]));
            builder.unique();

            IndexLabel label = builder.ifNotExist().create();

            indexLabels.add(label.name());
        }

    }
}
